#!/usr/bin/env python

import gi
import traceback
import os,subprocess
gi.require_version('OSTree', '1.0')
from gi.repository import GLib, Gio, OSTree,GObject
from SystemUpdater.Core.errors import *
from SystemUpdater.Core.enums import *
from gettext import gettext as _

RETRY_LIMIT = 5
ORIGIN_REMOTE_NAME = "openkylin"

class BackendBaseOstree(GObject.GObject):    
    __gsignals__ = {
                    "transaction-done": (GObject.SignalFlags.RUN_FIRST,
                                         None,
                                         (GObject.TYPE_PYOBJECT,)),
                    "download-info": (GObject.SIGNAL_RUN_FIRST,
                                GObject.TYPE_NONE,
                                (   GObject.TYPE_INT,
                                    GObject.TYPE_INT,
                                    GObject.TYPE_INT64,
                                    GObject.TYPE_INT64
                                ))
                    }
    
    def __init__(self,parent,sysroot_path=None,osname=None):
        GObject.GObject.__init__(self)
        self._async_progress = None
        self.error_code = ''
        self.exit = None
        if sysroot_path:
            sysroot_path = Gio.File.new_for_path(sysroot_path)
            
        self.sysroot = OSTree.Sysroot.new(sysroot_path)
        self.sysroot.set_mount_namespace_in_use()
        self.sysroot.initialize()
        self.sysroot.load()

        booted_deployment = self.sysroot.get_booted_deployment()
        if booted_deployment:
            self.osname = booted_deployment.get_osname()
        else:
            self.osname = osname

        self.merge_deployment = self.sysroot.get_merge_deployment(self.osname)
        if not self.merge_deployment:
            pass
        self.origin = self.merge_deployment.get_origin()

        # 异常情况下origin文件origin_remote为空 默认强制openkylin
        refspec = self.origin.get_string("origin","refspec")
        ret,origin_remote,origin_ref = OSTree.parse_refspec(refspec)
        if origin_remote == None:
            origin_refspec = ORIGIN_REMOTE_NAME+":"+origin_ref
            self.origin.set_string("origin","refspec",origin_refspec)

    def run(self,trans=None):
        logging.info("Processing transaction")
        GLib.idle_add(self._run_transaction_idle,trans)

    def _run_transaction_idle(self,trans):
        """Run the transaction"""
        try:
            self._async_progress = OSTree.AsyncProgress.new()
            self._async_progress.connect("changed",self._on_async_progress)
            self._try_lock()
            self._run_transaction(self.sysroot,trans)
        except GLib.Error as error:
            logging.error(str(error))
            self.exit = False
            self._make_error(error,trans)
        except UpdateBaseError as excep:
            self.exit = False
            self.error_code = excep.code
        except Exception as excep:
            logging.error(str(excep))
            self.exit = False

            tbk = traceback.format_exc()
            logging.error(str(tbk))
        else:
            self.exit = True
        
        finally:
            self._try_unlock()
            self._async_progress.finish()
            logging.info("Finished transaction")

            self._emit_transaction_done(trans)

        return False

    def _run_transaction(self,sysroot,trans):
        """This method needs to be implemented by the backends."""
        pass
        # raise errors.TransactionFailed(enums.ERROR_NOT_SUPPORTED)

    def _try_lock(self):
        ret,out_acquired = self.sysroot.try_lock()
        if not out_acquired:
            raise UpdateBaseError(ERROR_OTHER_PROGRESS_OCCUPATION)
    
    def _try_unlock(self):
        self.sysroot.unlock()

    def _emit_transaction_done(self,trans):
        """Emit the transaction-done signal.

        Keyword argument:
        trans -- the finished transaction
        """
        # logging.debug("Emitting transaction-done: %d",)
        self.emit("transaction-done",trans)

    def _on_async_progress(self,obj):
        try:
            start_time = obj.get_uint64("start-time")
            bytes_transferred = obj.get_uint64("bytes-transferred") 

            fetched = obj.get_uint("fetched")
            requested = obj.get_uint("requested")

            # 下载速度
            elapsed_secs=(GLib.get_monotonic_time()-start_time)/GLib.USEC_PER_SEC
            bytes_sec = int(bytes_transferred / elapsed_secs)

            if requested <= 2:
                return
            
            self.emit("download-info",fetched,requested,bytes_transferred,bytes_sec)
        except Exception as excep:
            tbk = traceback.format_exc()
            logging.error(str(tbk))

    def _make_error(self,error,trans):
        if error.code in (0,37):
            self.error_code = ERROR_NETWORK_FAILED
        #取消
        elif error.code == 19:
            trans.cancellable.reset()
            self.error_code = ERROR_CANCELLED
        elif error.code == 1:
            #Remote "kylin" not found (1)
            self.exit = True

class UpdateBackend(BackendBaseOstree):

    def __init__(self,parent,sysroot_path=None,osname=None):
        """Initialize a new AptWorker instance."""
        BackendBaseOstree.__init__(self,parent,sysroot_path,osname)

    def _run_transaction(self, sysroot,trans):
        ostree_repo = sysroot.repo()
        from_revision = self.merge_deployment.get_csum()
        
        refspec = self.origin.get_string("origin","refspec")
        ret,origin_remote,origin_ref = OSTree.parse_refspec(refspec)

        g_options = GLib.Variant("a{sv}",
                                {   "refs":GLib.Variant("as",[origin_ref]),\
                                    "n-network-retries":GLib.Variant("u",RETRY_LIMIT),\
                                    "flags":GLib.Variant("i",OSTree.RepoPullFlags.COMMIT_ONLY),\
                                    # "timestamp-check":GLib.Variant("b",True),\
                                    "depth":GLib.Variant("i",1)
                                })
        
        logging.info("Start pull metadata refs:%s...",origin_ref)
        ret = ostree_repo.pull_with_options(origin_remote,g_options,self._async_progress,trans.cancellable)

        origin_refspec = origin_remote+":"+origin_ref
        ret,new_revision = ostree_repo.resolve_rev(origin_refspec,True)
        new_ref = origin_ref

        ret,metadata = ostree_repo.load_variant(OSTree.ObjectType.COMMIT,new_revision)
        
        n_metadata = metadata[0]

        # 检查是否当前分支已经终止，切换到新的分支
        if OSTree.COMMIT_META_KEY_ENDOFLIFE_REBASE in n_metadata:
            new_ref = n_metadata[OSTree.COMMIT_META_KEY_ENDOFLIFE_REBASE]
            # 拉取新分支的元数据
            g_options = GLib.Variant("a{sv}",
                                    {   "refs":GLib.Variant("as",[new_ref]),\
                                        "n-network-retries":GLib.Variant("u",RETRY_LIMIT),\
                                        "flags":GLib.Variant("i",OSTree.RepoPullFlags.COMMIT_ONLY),\
                                        # "timestamp-check":GLib.Variant("b",True),\
                                        "depth":GLib.Variant("i",1)
                                    })
            logging.info("From origin ref:%s To new refs:%s...",origin_ref,new_ref)
            logging.info("Start pull new metadata refs:%s...",new_ref)
            ret = ostree_repo.pull_with_options(origin_remote,g_options,self._async_progress,trans.cancellable)

            origin_refspec = origin_remote+":"+new_ref
            ret,new_revision = ostree_repo.resolve_rev(origin_refspec,True)

        ret,metadata = ostree_repo.load_variant(OSTree.ObjectType.COMMIT,new_revision) 
        
        # 判断是否有新的更新
        trans.from_revision = from_revision
        trans.new_revision = new_revision
        trans.available_refs = new_ref

        trans.available_metadata = metadata[0]

class DownloadBackend(BackendBaseOstree):
    def __init__(self,parent,sysroot_path=None,osname=None):
        """Initialize a new AptWorker instance."""
        BackendBaseOstree.__init__(self,parent,sysroot_path,osname)

    def _run_transaction(self, sysroot,trans):
        ostree_repo = sysroot.repo()
        refspec = self.origin.get_string("origin","refspec")
        ret,origin_remote,origin_ref = OSTree.parse_refspec(refspec)

        logging.info("start pull data from available_refs:%s...",trans.available_refs)
        g_options = GLib.Variant("a{sv}",
                                 {
                                    "refs":GLib.Variant("as",[trans.available_refs]),\
                                    "n-network-retries":GLib.Variant("u",RETRY_LIMIT),\
                                    # "timestamp-check":GLib.Variant("b",True),\
                                    # "low-speed-limit-bytes":GLib.Variant("u",10),\
                                    # "low-speed-limit-seconds":GLib.Variant("u",20),\
                                    "flags":GLib.Variant("i",OSTree.RepoPullFlags.NONE),\
                                    "depth":GLib.Variant("i",1)
                                })

        ostree_repo.pull_with_options(origin_remote,g_options,self._async_progress,trans.cancellable)

        # 下载完成后检查下载的元数据的时间戳
        # upgrader.check_timestamps(ostree_repo,from_revision,trans.new_revision)

class DeployBackend(BackendBaseOstree):
    def __init__(self,parent,sysroot_path=None,osname=None):
        """Initialize a new AptWorker instance."""
        BackendBaseOstree.__init__(self,parent,sysroot_path,osname)

    def _deployment_sanitycheck_true(self,rootfs_path):
        try:
            pid = os.fork()
            if pid == 0:
                os.chroot(rootfs_path)
                os.chdir("/")
                # 执行命令
                result = subprocess.run(["/usr/bin/true"],capture_output=True, text=True)
                os._exit(result.returncode)
            else:
                # 等待子进程结束
                _, exit_code = os.waitpid(pid, 0)
                logging.info("Command exited with status: %d",os.WEXITSTATUS(exit_code))
                if os.WIFEXITED(exit_code) != 0:
                    logging.error("Command did not exit normally.")
        except Exception as e:
            logging.error(str(e))

    def _run_transaction(self, sysroot,trans):
        ostree_repo = sysroot.repo()

        new_revision = trans.new_revision
        refs = trans.available_refs

        # 获取当前origin文件
        refspec = self.origin.get_string("origin","refspec")
        ret,origin_remote,origin_ref = OSTree.parse_refspec(refspec)

        # 修改origin 为当前部署新分支
        origin_refspec = origin_remote+":"+refs
        self.origin.set_string("origin","refspec",origin_refspec)

        self._start_plymouth_splash()
        self._check_plymouth_state()

        self._message_to_plymouth(_("The system is installing updates. Please do not turn off your computer. "))
        self._progress_to_plymouth(10)

        #执行清理任何先前部分失败留下的数据。
        ret = sysroot.prepare_cleanup(None)

        self._progress_to_plymouth(60)
        
        logging.info("start deploy available_refs:%s new_revision:%s...",refs,new_revision)
        ret,upgrade_deployment = sysroot.deploy_tree(self.osname,new_revision,self.origin,self.merge_deployment,None,None)

        # 进行部署时 保留历史部署
        flags = OSTree.SysrootSimpleWriteDeploymentFlags.RETAIN
        sysroot.simple_write_deployment(self.osname,upgrade_deployment,self.merge_deployment,flags,None)

        self._progress_to_plymouth(90)

        special_deployment = []
        has_deployment = []
        old_deployments = sysroot.get_deployments()
        for deploys in old_deployments:
            # 放hotfix和pin 部署
            if deploys.get_unlocked() == OSTree.DeploymentUnlockedState.HOTFIX or deploys.is_pinned():
                special_deployment.append(deploys)
            # 其他升级保存的部署
            elif not upgrade_deployment.equal(deploys):
                has_deployment.append(deploys)

        logging.info("Number of current deployments:%d and all:%d",len(has_deployment),len(old_deployments))
        if len(has_deployment) > 5:
            logging.info("The number of deployments exceed 5 and delete the 5 after the depoyment...")
            new_deployment = [upgrade_deployment] + special_deployment + has_deployment[:5]
            ret = sysroot.write_deployments(new_deployment,None)
            self._progress_to_plymouth(95)

        # 健全性检查 
        # rootfs_path = sysroot.get_deployment_dirpath(new_deployment)
        # self._deployment_sanitycheck_true(rootfs_path)

        ret,metadata = ostree_repo.load_variant(OSTree.ObjectType.COMMIT,new_revision) 

        trans.new_revision = new_revision
        trans.available_metadata = metadata[0]

    def _start_plymouth_splash(self):
        try:
            logging.info("Running plymouth --splash")
            if os.path.exists("/sbin/plymouthd"):
                subprocess.run(["/sbin/plymouthd", "--mode=boot","--attach-to-session"],timeout=1)
            if os.path.exists("/bin/plymouth"):
                subprocess.Popen(["/bin/plymouth", "show-splash","--wait"])
                subprocess.call(["/bin/plymouth","system-update","--progress=0"])
        except Exception as exc:
            logging.error(exc)

    def _check_plymouth_state(self):
        try:
            args = ["/bin/plymouth","--ping"]
            p = subprocess.run(args,timeout=2,stdout=subprocess.PIPE,stderr=subprocess.STDOUT,text=True)
            if p.returncode == 0:
                logging.info("Check: Plymouthd state start to success...")
            else:
                logging.info("Check: Plymouthd Failed to Boot and Restart boot up again...")
                self._start_plymouth_splash()
        except Exception as exc:
            logging.error(exc) 

    def _message_to_plymouth(self,message):
        subprocess.call(["/bin/plymouth", "message", "--text", message])

    def _progress_to_plymouth(self,progress):
        tmp = ("--progress=%d"%progress)
        run_cmd = ["/bin/plymouth","system-update",tmp]
        subprocess.call(run_cmd)

class RollbackBackend(BackendBaseOstree):
    def __init__(self,parent,sysroot_path=None,osname=None):
        """Initialize a new AptWorker instance."""
        BackendBaseOstree.__init__(self,parent,sysroot_path,osname)

    def _run_transaction(self, sysroot,trans):
        new_deployments = []

        out_pending,out_rollback = sysroot.query_deployments_for(self.osname)

        if out_rollback == None:
            raise UpdateBaseError(ERROR_NOT_ROLLBAK_DEPLOYMENT)

        rollback_csum = out_rollback.get_csum()

        old_deployments = sysroot.get_deployments()

        new_deployments.append(out_rollback)

        # 部署的列表
        for deploys in old_deployments:
            if not out_rollback.equal(deploys):
                new_deployments.append(deploys)

        logging.info("start reollback checksum:%s...",rollback_csum)
        ret = sysroot.write_deployments(new_deployments,None)

        origin = out_rollback.get_origin()
        refspec = origin.get_string("origin","refspec")
        ret,origin_remote,origin_ref = OSTree.parse_refspec(refspec)

        trans.available_refs = origin_ref

class LoadCacheBackend(BackendBaseOstree):
    def __init__(self,parent,sysroot_path=None,osname=None):
        """Initialize a new AptWorker instance."""
        BackendBaseOstree.__init__(self,parent,sysroot_path,osname)

    def _run_transaction(self, sysroot,trans):
        trans.cur_refs = self._get_refs()
        trans.rollback_deployment = self._rollback_deployment(sysroot,self.osname)

    def _get_refs(self):
        refspec = self.origin.get_string("origin","refspec")
        ret,origin_remote,origin_ref = OSTree.parse_refspec(refspec)
        if origin_ref == None:
            origin_ref = ''
        return origin_ref
    
    def _rollback_deployment(self,sysroot,osname):
        out_pending,out_rollback = sysroot.query_deployments_for(osname)
        
        if out_rollback and not out_rollback.is_pinned():
            rollback_origin = out_rollback.get_origin()
            refspec = rollback_origin.get_string("origin","refspec")
            ret,origin_remote,origin_ref = OSTree.parse_refspec(refspec)
            
            return [origin_ref]
        else:
            return []